// test.js

import { GoogleGenerativeAI } from "@google/generative-ai";
import * as fs from "node:fs";

async function main() {
  // --- 配置区域 ---

  // 1. 设置您的 API 密钥
  // 推荐方式：在终端设置环境变量 API_KEY
  // 备用方式：直接在此处提供: const genAI = new GoogleGenerativeAI("YOUR_API_KEY");
  const genAI = new GoogleGenerativeAI(process.env.API_KEY);
  if (!process.env.API_KEY) {
    console.error("错误：请设置 API_KEY 环境变量。");
    return;
  }

  // 2. 定义输入和输出文件路径
  const inputImagePath = "rice.jpeg";
  const outputImagePath = "generated-figurine.png";

  // 3. 您的详细提示词
  const prompt = `create a 1/7 scale commercialized figurine of the characters in the picture, in a realistic style, in a real environment. The figurine is placed on a computer desk. The figurine has a round transparent acrylic base, with no text on the base. The content on the computer screen is the Zbrush modeling process of this figurine. Next to the computer screen is a BANDAI-style toy packaging box printed with the original artwork. The packaging features two-dimensional flat illustrations.`;

  // --- 脚本执行区域 ---

  // 检查输入文件是否存在
  if (!fs.existsSync(inputImagePath)) {
    console.error(`错误：找不到输入文件 ${inputImagePath}。请确保图片和脚本在同一个目录下。`);
    return;
  }

  // 将本地图像文件转换为 Google AI API 所需的格式
  function fileToGenerativePart(path, mimeType) {
    return {
      inlineData: {
        data: Buffer.from(fs.readFileSync(path)).toString("base64"),
        mimeType,
      },
    };
  }

  // 获取一个支持图像输入和输出的模型
  const model = genAI.getGenerativeModel({ model: "gemini-2.5-flash-image-preview" });

  console.log(`正在读取图片 ${inputImagePath} 并根据提示词生成新图片...`);
  console.log("这可能需要一些时间，请稍候。");

  try {
    // 准备要发送给模型的输入内容（提示词 + 图片）
    const imageParts = [
      fileToGenerativePart(inputImagePath, "image/jpeg"), // 假设 rice.jpg 是 jpeg 格式
    ];

    const result = await model.generateContent([prompt, ...imageParts]);
    const response = result.response;

    // 检查模型是否有返回内容
    if (!response.candidates || response.candidates.length === 0 || !response.candidates[0].content) {
        // 有时候模型会因为安全设置等原因拒绝生成，并返回文本提示
        const feedbackText = response.promptFeedback?.blockReason?.message || response.text();
        console.error("错误：模型没有生成有效的图片。可能的原因:", feedbackText || "未知错误，请检查您的提示词或输入图片。");
        return;
    }

    let imageGenerated = false;
    // 遍历返回的所有部分，寻找图像数据
    for (const part of response.candidates[0].content.parts) {
      if (part.inlineData) {
        // 这是图像数据！
        const imageData = part.inlineData.data; // 这是 base64 编码的字符串
        const mimeType = part.inlineData.mimeType;

        // 将 base64 字符串转换为二进制 Buffer
        const buffer = Buffer.from(imageData, "base64");

        // 将 Buffer 写入新文件
        fs.writeFileSync(outputImagePath, buffer);
        console.log(`\n🎉 图片已成功生成并保存为：${outputImagePath}`);
        imageGenerated = true;
        break; // 找到一张图片就够了，退出循环
      }
    }

    if (!imageGenerated) {
        console.error("错误：模型的响应中未找到图像数据。它可能返回了文本：", response.text());
    }

  } catch (error) {
    console.error("请求 Google AI 时发生严重错误:", error);
  }
}

main();